import os
from typing import List, AnyStr
import pymupdf as mu
from math import ceil, sqrt

CITE_STYLE = [
    '[]',
    '()',
    '【】'
]

class SentenceCache(object):

    def __init__(self, max_size=3):
        self._sents = []
        self._max = max_size
        self._wrds = ""
        self._scnt = 0
        self._wcnt = 0

    def push(self, next: AnyStr):
        if len(self._wrds) > 0:
            if self._wrds[-1] == '-':
                self._wrds = self._wrds[:-1]
            else:
                self._wrds += " "
                self._wcnt += 1
        self._wrds += next
        if next.find('i.e') < 0 and next.find('e.g') < 0 and next.find('w.r.t') < 0 and \
              next.find('Fig.') < 0 and next.find('al.') < 0 and next[-1] == '.' and \
              self._wcnt > 3:
            self._scnt += 1
            self._wcnt = 0
            self._sents.append(self._wrds)
            self._wrds = ""
            if len(self._sents) > self._max:
                self._sents.pop(0)
            return True
        return False
    
    @property
    def size(self):
        return self._max

    def get_sent_no(self):
        return self._scnt

    def dump(self):
        return self._sents


def extract_digit(clip_str: AnyStr):
    ret = []
    buf = ''
    digit_beg = False
    for c in clip_str:
        if c.isdigit():
            buf += c
            digit_beg = True
        else:
            if not digit_beg:
                pass
            else:
                digit_beg = False
                ret.append(int(buf))
                buf = ''
    return ret[0] if len(ret) == 1 else ret

def search_digit_cite(page_txt: AnyStr, label, cache_size = 3, cite_style=0):
    ret = []
    label = str(label)
    brac = CITE_STYLE[cite_style]
    cite_beg = None
    cite_end = None
    index = 0
    while index < len(page_txt):
        c = page_txt[index]
        if c == brac[0]:
            cite_beg = index
        elif c == brac[1] and cite_beg is not None:
            cite_end = index
            found_loc = page_txt[cite_beg:cite_end].find(label) 
            if found_loc >= 0 and page_txt[cite_beg + found_loc-1] in [brac[0], ',', ' ', '-', '\n'] and \
              page_txt[cite_beg + found_loc + len(label)] in [brac[1], ',', ' ', '-', '\n']:
                # print("found cite: " + page_txt[cite_beg:cite_end+1])
                sent_cnt = 0
                rbeg = cite_beg - 1
                rend = cite_end + 1
                while sent_cnt < cache_size // 2 and rbeg > 0:
                    if page_txt[rbeg] == '.' and not page_txt[rbeg+1].isspace() and \
                      page_txt[rbeg-2:rbeg] != 'al' and page_txt[rbeg-3:rbeg] != 'i.e' and \
                      page_txt[rbeg-3:rbeg] != 'e.g' and page_txt[rbeg-3:rbeg] != 'Fig' and \
                      page_txt[rbeg-5:rbeg] != 'w.r.t':
                        sent_cnt += 1
                    rbeg -= 1

                sent_cnt = 0
                while sent_cnt < ceil(cache_size/2) and rend < len(page_txt):
                    if page_txt[rend] == '.' and not page_txt[rend+1].isspace() and \
                      page_txt[rend-2:rend] != 'al' and page_txt[rend-3:rend] != 'i.e' and \
                      page_txt[rend-3:rend] != 'e.g' and page_txt[rend-3:rend] != 'Fig' and \
                      page_txt[rend-5:rend] != 'w.r.t':
                        sent_cnt += 1
                    rend += 1
                ret.append(page_txt[rbeg:rend].replace('\n',' '))
                index = rend + 1
                cite_beg = cite_end = None
        index += 1
    return ret

def read_page(page, cite, found_scnt, cache, c=None):
    blks = page.get_text('words', clip=c)
    for b in blks:
        sent_end = cache.push(b[4])
        brt = mu.Point(b[0:2])
        if brt.distance_to(cite.tl) < cite.height/2:
            found_scnt = cache.get_sent_no()
        elif found_scnt > 0 and sent_end and (cache.get_sent_no() - found_scnt) == ceil(cache.size/2):
            found_scnt = -1
            mention = cache.dump()
            return mention
    return found_scnt


def locate_cited_title(tar: AnyStr, doc: mu.Document, extra: List[AnyStr] = None):
    found_tl = None
    found_page = -1
    ref_line_h = 0
    for rp in reversed(doc):
        rl = rp.search_for(tar, quads=False)
        
        found_ext = []
        if extra is not None:
            for e in extra:
                found_ext += rp.search_for(e, quads=False)

        if len(rl) > 0:
            x = rl[0].x0
            y = rl[0].y0
            for r in rl:
                x = min(x, r.x0)
                y = min(y, r.y0)

            if len(found_ext) > 0:
                ex = found_ext[0].x0
                ey = found_ext[0].y0
                for ext in found_ext:
                    ex = min(ex, ext.x0)
                    ey = min(ey, ext.y0)
                if abs(ey - y) < 30 and abs(ex - x) < 50:
                    x = min(x, ex)
                    y = min(y, ey)

            if len(rl) > 1:
                ref_line_h = rl[1].y0 - rl[0].y0
            else:
                ref_line_h = rl[0].height
            
            y = rp.bound().y1 - y
            found_tl = mu.Point(x, y)
            found_page = rp.number
            break

    if found_tl is None:
        for rp in reversed(doc):
            found_tl, ref_line_h = search_for_fallback(rp, tar, extra)
            if found_tl is not None:
                # Debug
                # label_rect = mu.Rect(found_tl.x-20, 
                #                      rp.rect.height - found_tl.y, 
                #                      found_tl.x, 
                #                      rp.rect.height - found_tl.y + ref_line_h)
                # rp.add_circle_annot(label_rect)
                
                found_page = rp.number
                break

    return found_page, found_tl, ref_line_h


def search_for_fallback(page, tar: AnyStr, extra: List[AnyStr], allow_dup = False):
    tars = tar.lower().split(' ')
    cursor = 0
    lh = -.1
    matched_ws = []
    extra_ws = []
    found = False
    half = False
    
    for w in page.get_text('words'):
        if w[4].lower().find(tars[cursor]) >= 0:
            cursor += 1
            matched_ws.append(mu.Rect(w[0:4]))
            lh = max(lh, matched_ws[-1].height)
            found = True
            half = False
            if cursor == len(tars):
                if allow_dup:
                    cursor = 0
                else:
                    break
        elif tars[cursor].find(w[4].lower()) == 0:
            matched_ws.append(mu.Rect(w[0:4]))
            lh = max(lh, matched_ws[-1].height)
            half = True
        elif half and tars[cursor].find(w[4].lower()) > 0:
            matched_ws.append(mu.Rect(w[0:4]))
            lh = max(lh, matched_ws[-1].height)
            if w[4][-1].lower() == tars[cursor][-1]:
                cursor += 1
        elif found or half:
            found = False
            half = False
            cursor = 0
            matched_ws.clear()

    if extra is not None:
        for e in extra:
            found_extra = search_for_fallback(page, e, None, True)
            if isinstance(found_extra, List):
                extra_ws += found_extra

    if found:
        if allow_dup:
            found_dups = []
            for mw in matched_ws:
                found_dups.append(mu.Point(mw.x0, page.rect.height - mw.y0))
            return found_dups
        else:
            left, top, el, et = None, None, None, None
            left = matched_ws[0].x0
            top = matched_ws[0].y0
            for mw in matched_ws:
                left = min(left, mw.x0) 
                top = min(top, mw.y0) 
            # reverse the top coordinate because et has been reversed   
            top = page.rect.height-top

            if len(extra_ws) > 0:
                min_dist = 999.
                el = left
                et = top
                for mw in extra_ws:
                    cur_dist = sqrt((mw.x-left)**2 + 100*(mw.y-top)**2)
                    if min_dist >= cur_dist:
                        el = mw.x
                        et = mw.y
                        min_dist = cur_dist
        
                left = min(left, el)
                top = max(top, et)

            return mu.Point(left, top) if top is not None else None, lh
    else:
        return None, -1
    
        

def search_pdf(fileName, targetName, cache_size = 3, clip = (0.1, 0.9), extra_name = None):
    
    doc = mu.open(fileName)
    found_page, found_tl, Line_Height = locate_cited_title(targetName, doc, extra_name)
    if found_tl is None:
        return {}

    ret = {}

    def _append_result(mention: List[AnyStr]):
        if 'cite' not in ret.keys():
            ret['cite'] = []
        cite = ''
        for s in mention:
            cite += s + ' '
        ret['cite'].append(cite)

    page_clip = doc[0].rect
    page_clip.y0 = page_clip.y1*clip[0]
    page_clip.y1 = page_clip.y1*clip[1]
    page_clip = page_clip.irect

    cache = None
    found_scnt = -1
    last_frm = None

    for page in doc:
        if page.number == found_page:
            break

        if found_scnt > 0:
            mention = read_page(page, last_frm, found_scnt, cache, page_clip)
            assert(isinstance(mention, List))
            _append_result(mention)
            last_frm = None
            found_scnt = -1
            cache = SentenceCache(cache_size)

        for link in page.links(kinds=[4]):
            if 'to' in link.keys() and link['page'] == found_page:
                frm = link['from']
                to = link['to']
                vdist = abs(to.y - found_tl.y)
                hdist = abs(to.x - found_tl.x)
                if vdist < 1.5*Line_Height and hdist < 0.25*page.rect.width:
                    if 'label' not in ret.keys():
                        clips = page.get_textbox(link['from'])
                        ret['label'] = clips
                    # debug
                    tmp_rect = mu.Point(to.x, page.rect.height - to.y)
                    doc[link['page']].add_text_annot(tmp_rect, page.get_textbox(frm))

                    cache = SentenceCache(cache_size)
                    mention = read_page(page, frm, found_scnt, cache, page_clip)
                    if isinstance(mention, List):
                        _append_result(mention)
                    else:
                        found_scnt = mention
                        last_frm = frm
        
        if page.load_links() is None:
            label_rect = mu.Rect(found_tl.x-20, 
                                 page.rect.height - found_tl.y, 
                                 found_tl.x, 
                                 page.rect.height - found_tl.y + Line_Height)
            
            ret['label'] = doc[found_page].get_textbox(label_rect).strip()
            label = extract_digit(ret['label'])
            page_txt = page.get_text('text', clip=page_clip)
            found_cites = search_digit_cite(page_txt, label, cache_size)
            if 'cite' not in ret.keys():
                ret['cite'] = []
            ret['cite'] += found_cites
    
    doc.save(os.path.join('out', fileName.split(os.path.sep)[-1] + '-tmp.pdf'))
    return ret


if __name__ == '__main__':
    tar = 'SMAT: An input adaptive auto-tuner for sparse matrix-vector multiplication'
    extra = [
        'J. Li',
        'Li J.',
        'Jiajia Li',
        'Li Jiajia',
    ]
    prefix = '../161-170'
    # Debug
    # ret = search_pdf('test.pdf', tar, extra_name=extra)
    # print(ret)
    # exit()
    for pdf in os.listdir(prefix):
        if pdf[-4:] != '.pdf':
            continue
        ret = search_pdf(os.path.join(prefix, pdf), tar, cache_size=3, extra_name=extra)
        with open(os.path.join('tmp', pdf+'.txt'), 'w') as f:
            lines = []
            lines.append('\n'+'--'*10+'\n\n')
            lines.append('- 引用文献： ' + pdf + '\n')
            try:
                lines.append('- 引用标签：' + ret['label'] + '\n')
                lines.append('- 引用：' + '\n')
                for cite in ret['cite']:
                    lines.append("   - " + cite + '\n')
            except KeyError:
                lines.append('**failed to load cross-refs**' + '\n')
            lines.append('\n'+'--'*10+'\n\n')
            f.writelines(lines)
            


